#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import os
import time
import random
import threading
import psutil
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QLabel, QComboBox, QSpinBox, QProgressBar,
    QMessageBox, QGroupBox, QGridLayout, QFileDialog, QTabWidget,
    QFrame, QSizePolicy, QDoubleSpinBox
)
from PySide6.QtCore import Qt, QThread, Signal, Slot
from PySide6.QtGui import QIcon, QFont, QColor, QPalette


class DiskTestThread(QThread):
    """磁盘测试线程类，用于在后台执行磁盘读写测试"""
    
    # 定义信号
    update_progress = Signal(int)
    update_write_speed = Signal(float)
    update_read_speed = Signal(float)
    test_complete = Signal(dict)
    test_error = Signal(str)
    
    def __init__(self, test_dir, file_size_mb, block_size_kb, test_type):
        super().__init__()
        self.test_dir = test_dir
        self.file_size_mb = file_size_mb
        self.block_size_kb = block_size_kb
        self.test_type = test_type  # 'write', 'read', or 'both'
        self.running = True
        
    def run(self):
        try:
            results = {}
            
            # 测试文件名
            test_file = os.path.join(self.test_dir, f"disk_speed_test_{int(time.time())}.tmp")
            
            # 计算总字节数和块大小
            total_bytes = int(self.file_size_mb * 1024 * 1024)
            block_size = int(self.block_size_kb * 1024)
            
            # 生成随机数据块
            random_data = bytearray(random.getrandbits(8) for _ in range(block_size))
            
            # 如果是读取测试且文件不存在，先创建文件
            if self.test_type == 'read' and not os.path.exists(test_file):
                # 预先写入文件用于读取测试
                self._write_test_file(test_file, total_bytes, block_size, random_data)
            
            # 执行写入测试
            if self.test_type in ['write', 'both']:
                write_speed = self._test_write(test_file, total_bytes, block_size, random_data)
                results['write_speed'] = write_speed
                if not self.running:
                    self._cleanup(test_file)
                    return
            
            # 执行读取测试
            if self.test_type in ['read', 'both'] and os.path.exists(test_file):
                read_speed = self._test_read(test_file, total_bytes, block_size)
                results['read_speed'] = read_speed
                if not self.running:
                    self._cleanup(test_file)
                    return
            
            # 清理临时文件
            self._cleanup(test_file)
            
            # 发送完成信号
            self.test_complete.emit(results)
            
        except Exception as e:
            self.test_error.emit(str(e))
    
    def _test_write(self, test_file, total_bytes, block_size, random_data):
        """执行写入测试"""
        bytes_written = 0
        start_time = time.time()
        
        try:
            with open(test_file, 'wb') as f:
                while bytes_written < total_bytes and self.running:
                    # 计算本次写入大小
                    remaining = total_bytes - bytes_written
                    write_size = min(block_size, remaining)
                    
                    # 写入数据
                    f.write(random_data[:write_size])
                    f.flush()
                    os.fsync(f.fileno())
                    
                    bytes_written += write_size
                    
                    # 计算进度和速度
                    progress = int((bytes_written / total_bytes) * 50)  # 写入测试占50%
                    elapsed = time.time() - start_time
                    if elapsed > 0:
                        speed_mbps = (bytes_written / (1024 * 1024)) / elapsed
                        self.update_write_speed.emit(speed_mbps)
                    
                    self.update_progress.emit(progress)
            
            # 最终速度
            elapsed = time.time() - start_time
            return (bytes_written / (1024 * 1024)) / elapsed
            
        except Exception as e:
            self.test_error.emit(f"写入测试失败: {str(e)}")
            return 0
    
    def _test_read(self, test_file, total_bytes, block_size):
        """执行读取测试"""
        bytes_read = 0
        start_time = time.time()
        
        try:
            with open(test_file, 'rb') as f:
                while bytes_read < total_bytes and self.running:
                    # 计算本次读取大小
                    remaining = total_bytes - bytes_read
                    read_size = min(block_size, remaining)
                    
                    # 读取数据
                    _ = f.read(read_size)
                    
                    bytes_read += read_size
                    
                    # 计算进度和速度
                    progress = 50 + int((bytes_read / total_bytes) * 50)  # 读取测试占50%
                    elapsed = time.time() - start_time
                    if elapsed > 0:
                        speed_mbps = (bytes_read / (1024 * 1024)) / elapsed
                        self.update_read_speed.emit(speed_mbps)
                    
                    self.update_progress.emit(progress)
            
            # 最终速度
            elapsed = time.time() - start_time
            return (bytes_read / (1024 * 1024)) / elapsed
            
        except Exception as e:
            self.test_error.emit(f"读取测试失败: {str(e)}")
            return 0
    
    def _cleanup(self, test_file):
        """清理临时文件"""
        try:
            if os.path.exists(test_file):
                os.remove(test_file)
        except:
            pass
    
    def stop(self):
        """停止测试"""
        self.running = False
        self.wait()


class DiskInfoWidget(QWidget):
    """磁盘信息显示组件"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.init_ui()
        self.update_disk_info()
    
    def init_ui(self):
        # 设置布局
        layout = QVBoxLayout(self)
        
        # 创建分组框
        disk_info_group = QGroupBox("磁盘信息")
        disk_info_layout = QGridLayout(disk_info_group)
        
        # 创建标签
        self.disk_labels = {
            "磁盘": [],
            "总容量": [],
            "已用": [],
            "可用": [],
            "使用率": []
        }
        
        # 设置标签样式
        font = QFont()
        font.setBold(True)
        
        # 添加表头
        for col, label_text in enumerate(self.disk_labels.keys()):
            header_label = QLabel(label_text)
            header_label.setFont(font)
            header_label.setAlignment(Qt.AlignCenter)
            disk_info_layout.addWidget(header_label, 0, col)
        
        # 添加刷新按钮
        refresh_button = QPushButton("刷新")
        refresh_button.clicked.connect(self.update_disk_info)
        layout.addWidget(refresh_button)
        
        layout.addWidget(disk_info_group)
        self.setLayout(layout)
    
    def update_disk_info(self):
        """更新磁盘信息"""
        # 清除现有标签
        for label_list in self.disk_labels.values():
            for label in label_list:
                label.deleteLater()
            label_list.clear()
        
        # 获取磁盘信息
        partitions = psutil.disk_partitions()
        
        # 为每个分区添加信息
        for row, partition in enumerate(partitions):
            try:
                # 获取分区使用情况
                usage = psutil.disk_usage(partition.mountpoint)
                
                # 计算容量
                total_gb = usage.total / (1024 ** 3)
                used_gb = usage.used / (1024 ** 3)
                free_gb = usage.free / (1024 ** 3)
                
                # 创建并添加标签
                self.disk_labels["磁盘"].append(QLabel(f"{partition.mountpoint} ({partition.device})"))
                self.disk_labels["总容量"].append(QLabel(f"{total_gb:.2f} GB"))
                self.disk_labels["已用"].append(QLabel(f"{used_gb:.2f} GB"))
                self.disk_labels["可用"].append(QLabel(f"{free_gb:.2f} GB"))
                
                # 创建进度条显示使用率
                usage_bar = QProgressBar()
                usage_bar.setValue(usage.percent)
                usage_bar.setFormat(f"{usage.percent}%")
                
                # 根据使用率设置颜色
                if usage.percent > 90:
                    usage_bar.setStyleSheet("QProgressBar { background-color: #f0f0f0; border: 1px solid #ccc; border-radius: 5px; } QProgressBar::chunk { background-color: #ff4444; }")
                elif usage.percent > 70:
                    usage_bar.setStyleSheet("QProgressBar { background-color: #f0f0f0; border: 1px solid #ccc; border-radius: 5px; } QProgressBar::chunk { background-color: #ffaa00; }")
                else:
                    usage_bar.setStyleSheet("QProgressBar { background-color: #f0f0f0; border: 1px solid #ccc; border-radius: 5px; } QProgressBar::chunk { background-color: #44ff44; }")
                
                # 将标签添加到布局
                disk_info_layout = self.findChild(QGridLayout)
                disk_info_layout.addWidget(self.disk_labels["磁盘"][-1], row + 1, 0)
                disk_info_layout.addWidget(self.disk_labels["总容量"][-1], row + 1, 1)
                disk_info_layout.addWidget(self.disk_labels["已用"][-1], row + 1, 2)
                disk_info_layout.addWidget(self.disk_labels["可用"][-1], row + 1, 3)
                disk_info_layout.addWidget(usage_bar, row + 1, 4)
                
            except Exception:
                # 跳过无法访问的分区
                continue


class DiskSpeedTester(QMainWindow):
    """磁盘速度测试主窗口"""
    
    def __init__(self):
        super().__init__()
        self.disk_test_thread = None
        self.init_ui()
    
    def init_ui(self):
        # 设置窗口标题和大小
        self.setWindowTitle("磁盘速度测试工具")
        self.resize(800, 600)
        
        # 创建中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建主布局
        main_layout = QVBoxLayout(central_widget)
        
        # 创建标签页
        tab_widget = QTabWidget()
        
        # 创建测试标签页
        test_tab = QWidget()
        test_layout = QVBoxLayout(test_tab)
        
        # 创建测试设置组
        settings_group = QGroupBox("测试设置")
        settings_layout = QGridLayout(settings_group)
        
        # 添加测试设置控件
        settings_layout.addWidget(QLabel("选择磁盘:"), 0, 0)
        self.disk_combo = QComboBox()
        self.update_disk_list()
        settings_layout.addWidget(self.disk_combo, 0, 1)
        
        # 自定义路径按钮
        custom_path_button = QPushButton("自定义路径")
        custom_path_button.clicked.connect(self.select_custom_path)
        settings_layout.addWidget(custom_path_button, 0, 2)
        
        settings_layout.addWidget(QLabel("测试文件大小 (MB):"), 1, 0)
        self.file_size_spin = QSpinBox()
        self.file_size_spin.setRange(10, 10240)
        self.file_size_spin.setValue(100)
        self.file_size_spin.setSingleStep(10)
        settings_layout.addWidget(self.file_size_spin, 1, 1)
        
        settings_layout.addWidget(QLabel("块大小 (KB):"), 2, 0)
        self.block_size_spin = QSpinBox()
        self.block_size_spin.setRange(4, 65536)
        self.block_size_spin.setValue(4096)
        self.block_size_spin.setSingleStep(4)
        settings_layout.addWidget(self.block_size_spin, 2, 1)
        
        settings_layout.addWidget(QLabel("测试类型:"), 3, 0)
        self.test_type_combo = QComboBox()
        self.test_type_combo.addItems(["写入测试", "读取测试", "读写测试"])
        settings_layout.addWidget(self.test_type_combo, 3, 1)
        
        # 创建测试控制组
        control_group = QGroupBox("测试控制")
        control_layout = QHBoxLayout(control_group)
        
        self.start_button = QPushButton("开始测试")
        self.start_button.clicked.connect(self.start_test)
        control_layout.addWidget(self.start_button)
        
        self.stop_button = QPushButton("停止测试")
        self.stop_button.clicked.connect(self.stop_test)
        self.stop_button.setEnabled(False)
        control_layout.addWidget(self.stop_button)
        
        # 创建进度显示组
        progress_group = QGroupBox("测试进度")
        progress_layout = QVBoxLayout(progress_group)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        progress_layout.addWidget(self.progress_bar)
        
        # 速度显示
        speed_layout = QHBoxLayout()
        
        self.write_speed_label = QLabel("写入速度: 0.00 MB/s")
        self.read_speed_label = QLabel("读取速度: 0.00 MB/s")
        
        speed_layout.addWidget(self.write_speed_label)
        speed_layout.addWidget(self.read_speed_label)
        
        progress_layout.addLayout(speed_layout)
        
        # 历史结果组
        results_group = QGroupBox("测试结果")
        results_layout = QVBoxLayout(results_group)
        
        self.results_text = QLabel("准备就绪")
        self.results_text.setWordWrap(True)
        results_layout.addWidget(self.results_text)
        
        # 添加到测试标签页
        test_layout.addWidget(settings_group)
        test_layout.addWidget(control_group)
        test_layout.addWidget(progress_group)
        test_layout.addWidget(results_group)
        
        # 创建磁盘信息标签页
        disk_info_tab = DiskInfoWidget()
        
        # 添加标签页
        tab_widget.addTab(test_tab, "速度测试")
        tab_widget.addTab(disk_info_tab, "磁盘信息")
        
        # 添加到主布局
        main_layout.addWidget(tab_widget)
        
        # 设置样式
        self.set_style()
    
    def set_style(self):
        """设置应用样式"""
        # 使用系统样式
        app.setStyle("Fusion")
        
        # 设置颜色主题
        palette = QPalette()
        palette.setColor(QPalette.Window, QColor(240, 240, 240))
        palette.setColor(QPalette.WindowText, Qt.black)
        palette.setColor(QPalette.Base, QColor(255, 255, 255))
        palette.setColor(QPalette.AlternateBase, QColor(240, 240, 240))
        palette.setColor(QPalette.ToolTipBase, Qt.white)
        palette.setColor(QPalette.ToolTipText, Qt.black)
        palette.setColor(QPalette.Text, Qt.black)
        palette.setColor(QPalette.Button, QColor(240, 240, 240))
        palette.setColor(QPalette.ButtonText, Qt.black)
        palette.setColor(QPalette.BrightText, Qt.red)
        palette.setColor(QPalette.Link, QColor(42, 130, 218))
        palette.setColor(QPalette.Highlight, QColor(42, 130, 218))
        palette.setColor(QPalette.HighlightedText, Qt.white)
        app.setPalette(palette)
        
        # 设置按钮样式
        self.setStyleSheet(""
            "QPushButton { "
            "    background-color: #4285F4; "
            "    color: white; "
            "    border-radius: 5px; "
            "    padding: 8px 16px; "
            "    font-weight: bold; "
            "}"
            "QPushButton:hover { "
            "    background-color: #5294FF; "
            "}"
            "QPushButton:pressed { "
            "    background-color: #3367D6; "
            "}"
            "QPushButton:disabled { "
            "    background-color: #CCCCCC; "
            "    color: #888888; "
            "}"
            "QGroupBox { "
            "    border: 1px solid #CCCCCC; "
            "    border-radius: 5px; "
            "    margin-top: 10px; "
            "    padding: 10px; "
            "}"
            "QGroupBox::title { "
            "    subcontrol-origin: margin; "
            "    left: 10px; "
            "    padding: 0 5px 0 5px; "
            "    background-color: transparent; "
            "}"
            "QProgressBar { "
            "    background-color: #f0f0f0; "
            "    border: 1px solid #ccc; "
            "    border-radius: 5px; "
            "    text-align: center; "
            "    height: 20px; "
            "}"
            "QProgressBar::chunk { "
            "    background-color: #4285F4; "
            "    border-radius: 4px; "
            "}"
            "QComboBox, QSpinBox { "
            "    padding: 5px; "
            "    border: 1px solid #CCCCCC; "
            "    border-radius: 4px; "
            "}"
        )
    
    def update_disk_list(self):
        """更新磁盘列表"""
        self.disk_combo.clear()
        partitions = psutil.disk_partitions()
        for partition in partitions:
            self.disk_combo.addItem(f"{partition.mountpoint} ({partition.device})")
    
    def select_custom_path(self):
        """选择自定义测试路径"""
        path = QFileDialog.getExistingDirectory(self, "选择测试目录")
        if path:
            self.disk_combo.addItem(path)
            self.disk_combo.setCurrentText(path)
    
    def start_test(self):
        """开始测试"""
        # 禁用开始按钮，启用停止按钮
        self.start_button.setEnabled(False)
        self.stop_button.setEnabled(True)
        self.progress_bar.setValue(0)
        self.results_text.setText("测试开始...")
        self.write_speed_label.setText("写入速度: 0.00 MB/s")
        self.read_speed_label.setText("读取速度: 0.00 MB/s")
        
        # 获取测试参数
        selected_disk = self.disk_combo.currentText().split(" ")[0]
        file_size_mb = self.file_size_spin.value()
        block_size_kb = self.block_size_spin.value()
        
        # 确定测试类型
        test_type_index = self.test_type_combo.currentIndex()
        test_type = ['write', 'read', 'both'][test_type_index]
        
        # 创建并启动测试线程
        self.disk_test_thread = DiskTestThread(
            selected_disk, file_size_mb, block_size_kb, test_type
        )
        
        # 连接信号
        self.disk_test_thread.update_progress.connect(self.update_progress)
        self.disk_test_thread.update_write_speed.connect(self.update_write_speed)
        self.disk_test_thread.update_read_speed.connect(self.update_read_speed)
        self.disk_test_thread.test_complete.connect(self.test_complete)
        self.disk_test_thread.test_error.connect(self.test_error)
        
        # 启动线程
        self.disk_test_thread.start()
    
    def stop_test(self):
        """停止测试"""
        if self.disk_test_thread:
            self.disk_test_thread.stop()
            self.disk_test_thread = None
        
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.results_text.setText("测试已停止")
    
    @Slot(int)
    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)
    
    @Slot(float)
    def update_write_speed(self, speed):
        """更新写入速度显示"""
        self.write_speed_label.setText(f"写入速度: {speed:.2f} MB/s")
    
    @Slot(float)
    def update_read_speed(self, speed):
        """更新读取速度显示"""
        self.read_speed_label.setText(f"读取速度: {speed:.2f} MB/s")
    
    @Slot(dict)
    def test_complete(self, results):
        """测试完成处理"""
        # 构建结果文本
        result_text = "测试完成！\n\n"
        
        if 'write_speed' in results:
            result_text += f"写入速度: {results['write_speed']:.2f} MB/s\n"
        
        if 'read_speed' in results:
            result_text += f"读取速度: {results['read_speed']:.2f} MB/s\n"
        
        # 重置界面状态
        self.results_text.setText(result_text)
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.disk_test_thread = None
    
    @Slot(str)
    def test_error(self, error_message):
        """测试错误处理"""
        QMessageBox.critical(self, "错误", f"测试过程中发生错误:\n{error_message}")
        self.results_text.setText(f"测试失败: {error_message}")
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.disk_test_thread = None
    
    def closeEvent(self, event):
        """关闭窗口时停止测试线程"""
        if self.disk_test_thread:
            self.disk_test_thread.stop()
        event.accept()


if __name__ == "__main__":
    # 创建应用实例
    app = QApplication(sys.argv)
    
    # 设置中文字体支持
    font = QFont("SimHei")
    app.setFont(font)
    
    # 创建并显示主窗口
    window = DiskSpeedTester()
    window.show()
    
    # 运行应用
    sys.exit(app.exec())